"""文档检索服务"""
import numpy as np
from typing import List, Dict, Any, Tuple
from app.services.storage import StorageService
from app.core.config import settings


class RetrievalService:
    """文档检索服务"""
    
    def __init__(self):
        """初始化检索服务"""
        self.storage_service = StorageService()
        self.top_k = settings.TOP_K
    
    def calculate_similarity(
        self,
        query_embedding: List[float],
        document_embeddings: List[List[float]]
    ) -> List[float]:
        """
        计算查询向量与文档向量的相似度（点积）
        
        Args:
            query_embedding: 查询向量
            document_embeddings: 文档向量列表
            
        Returns:
            相似度列表
        """
        query_vec = np.array(query_embedding)
        doc_vecs = np.array(document_embeddings)
        
        # 计算点积
        similarities = np.dot(doc_vecs, query_vec)
        
        return similarities.tolist()
    
    async def retrieve_top_chunks(
        self,
        query_embedding: List[float],
        top_k: int = None
    ) -> Tuple[List[Dict[str, Any]], List[float]]:
        """
        全局检索最相关的文本块
        
        Args:
            query_embedding: 查询向量
            top_k: 返回的文本块数量，如果不提供则使用默认值
            
        Returns:
            (最相关的文本块列表, 对应的相似度列表)
        """
        if top_k is None:
            top_k = self.top_k
        
        # 获取所有文本块（全局检索）
        chunks = await self.storage_service.get_all_chunks()
        
        if not chunks:
            return [], []
        
        # 提取嵌入向量
        embeddings = [chunk["embedding"] for chunk in chunks]
        
        # 计算相似度
        similarities = self.calculate_similarity(query_embedding, embeddings)
        
        # 获取 top_k 个最相似的索引
        top_indices = np.argsort(similarities)[-top_k:][::-1]
        
        # 提取对应的文本块和相似度
        top_chunks = [chunks[i] for i in top_indices]
        top_similarities = [similarities[i] for i in top_indices]
        
        return top_chunks, top_similarities

